import torch
from torch import nn
from backbones.image_transform import ImageTransform


class ImageClassifierModel(nn.Module):
    def __init__(self, embedding_dims, nhead, num_layers, num_classes=1000):
        super().__init__()
        self.embedding_dims = embedding_dims

        self.image_transform = ImageTransform(embedding_dims=embedding_dims)

        self.encoder_layer = nn.TransformerEncoderLayer(d_model=embedding_dims, nhead=nhead, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)

        self.fc = nn.Sequential(
            nn.Linear(embedding_dims, embedding_dims // 2),
            nn.ReLU(),
            nn.Linear(embedding_dims // 2, num_classes),
        )

    def forward(self, x):
        x = self.image_transform(x)
        x = self.transformer_encoder(x)
        x = self.fc(x[:, 0, :])  # 取出CLS所在的位置
        return x


if __name__ == '__main__':
    inputs = torch.randn(10, 3, 224, 224)
    model = ImageClassifierModel(1024, 8, 2, 3)
    outputs = model(inputs)
    print(outputs.shape)
